# -*- encoding: utf-8 -*-
import time
import cv2
import numpy as np
import matplotlib.pyplot as plt

# 1.加载可以识别物体的名称，将其存放在LABELS中，一共有80种，在这我们只使用car
labelsPath = "./yolo-coco/coco.names"
LABELS = open(labelsPath).read().strip().split("\n")
# 设置随机数种子，生成多种不同的颜色，当一个画面中有多个目标时，使用不同颜色的框将其框起来
np.random.seed(42)
COLORS = np.random.randint(0, 255, size=(200, 3),dtype="uint8")

# 加载已训练好的yolov3网络的权重和相应的配置数据
weightsPath = "./yolo-coco/yolov3.weights"
configPath = "./yolo-coco/yolov3.cfg"


# 加载好数据之后，开始利用上述数据恢复yolo神经网络
net = cv2.dnn.readNetFromDarknet(configPath, weightsPath)
# 获取YOLO中每一网络层的名称：['conv_0', 'bn_0', 'relu_0', 'conv_1', 'bn_1', 'relu_1', 'conv_2', 'bn_2', 'relu_2'...]
ln = net.getLayerNames()
# 获取输出层在网络中的索引位置，并以列表的形式：['yolo_82', 'yolo_94', 'yolo_106']
ln = [ln[i[0] - 1] for i in net.getUnconnectedOutLayers()]

# 2. 读取图像
frame = cv2.imread("./images/car2.jpg")
# 视频的宽度和高度，即帧尺寸
(W, H) = (None, None)
if W is None or H is None:
    (H, W) = frame.shape[:2]

# 根据输入图像构造blob,利用OPenCV进行深度网路的计算时，一般将图像转换为blob形式，对图片进行预处理，包括缩放，减均值，通道交换等
# 还可以设置尺寸，一般设置为在进行网络训练时的图像的大小
blob = cv2.dnn.blobFromImage(frame, 1 / 255.0, (416, 416), swapRB=True, crop=False)


# 3.将blob输入到前向网络中，并进行预测
net.setInput(blob)
start = time.time()
# yolo前馈计算，获取边界和相应的概率
# 输出layerOutsputs介绍：
# 是YOLO算法在图片中检测到的bbx的信息
# 由于YOLO v3有三个输出，也就是上面提到的['yolo_82', 'yolo_94', 'yolo_106']
# 因此layerOutsputs是一个长度为3的列表
# 其中，列表中每一个元素的维度是(num_detection, 85)
# num_detections表示该层输出检测到bbx的个数
# 85：因为该模型在COCO数据集上训练，[5:]表示类别概率；[0:4]表示bbx的位置信息；[5]表示置信度
layerOutputs = net.forward(ln)

# 下面对网络输出的bbx进行检查：
# 判定每一个bbx的置信度是否足够的高，以及执行NMS算法去除冗余的bbx
boxes = []  # 用于存放识别物体的框的信息，包括框的左上角横坐标x和纵坐标y以及框的高h和宽w
confidences = []  # 表示识别目标是某种物体的可信度
classIDs = []  # 表示识别的目标归属于哪一类，['person', 'bicycle', 'car', 'motorbike'....]

# 4. 遍历每一个输出层的输出
for output in layerOutputs:
    # 遍历某个输出层中的每一个目标
    for detection in output:
        scores = detection[5:]  # 当前目标属于某一类别的概率

        classID = np.argmax(scores)  # 目标的类别ID
        confidence = scores[classID]  # 得到目标属于该类别的置信度

        # 只保留置信度大于0.3的边界框，若图片质量较差，可以将置信度调低一点
        if confidence > 0.3:
            # 将边界框的坐标还原至与原图片匹配，YOLO返回的是边界框的中心坐标以及边界框的宽度和高度
            box = detection[0:4] * np.array([W, H, W, H])
            (centerX, centerY, width, height) = box.astype("int") # 使用 astype("int") 对上述 array 进行强制类型转换，centerX：框的中心点横坐标， centerY：框的中心点纵坐标，width：框的宽度，height：框的高度

            x = int(centerX - (width / 2))  # 计算边界框的左上角的横坐标
            y = int(centerY - (height / 2))  # 计算边界框的左上角的纵坐标

            # 更新检测到的目标框，置信度和类别ID
            boxes.append([x, y, int(width), int(height)])  # 将边框的信息添加到列表boxes
            confidences.append(float(confidence))  # 将识别出是某种物体的置信度添加到列表confidences
            classIDs.append(classID) # 将识别物体归属于哪一类的信息添加到列表classIDs


# 5. 非极大值抑制
idxs = cv2.dnn.NMSBoxes(boxes, confidences, 0.5, 0.3)


# 6. 获得最终的检测结果
dets = []  # 存放检测框的信息，包括左上角横坐标，纵坐标，右下角横坐标，纵坐标，以及检测到的物体的置信度,用于目标跟踪
if len(idxs) > 0:  # 存在检测框的话（即检测框个数大于0）
    for i in idxs.flatten():  #  循环检测出的每一个box
        # yolo模型可以识别很多目标，因为我们在这里只是识别车，所以只有目标是车的我们进行检测，其他的忽略
        if LABELS[classIDs[i]] == "car":
            (x, y) = (boxes[i][0], boxes[i][1])  # 得到检测框的左上角坐标
            (w, h) = (boxes[i][2], boxes[i][3])  # 得到检测框的宽和高
            cv2.rectangle(frame, (x, y), (x+w, y+h), (0,255,0), 2)  # 将方框绘制在画面上
            dets.append([x, y, x + w, y + h, confidences[i]])  # 将检测框的信息的放入dets中
# 设置数据类型，将整型数据转换为浮点数类型，且保留小数点后三位
np.set_printoptions(formatter={'float': lambda x: "{0:0.3f}".format(x)})
# 将检测框数据转换为ndarray,其数据类型为浮点型
dets = np.asarray(dets)

plt.imshow(frame[:,:,::-1])
plt.show()
